package annexb

import (
	"bytes"
	"encoding/binary"
	"encoding/hex"
	"strings"
	"testing"

	"github.com/stretchr/testify/require"
)

func decode(s string) []byte {
	b, _ := hex.DecodeString(strings.ReplaceAll(s, " ", ""))
	return b
}

func naluTypes(avcc []byte) (types []byte) {
	for {
		types = append(types, avcc[4])

		size := 4 + binary.BigEndian.Uint32(avcc)
		if size < uint32(len(avcc)) {
			avcc = avcc[size:]
		} else {
			break
		}
	}
	return
}

func TestFFmpegH264(t *testing.T) {
	// ffmpeg -re -i bbb.mp4 -c copy -f h264 -
	s := "000000016764001fac2484014016ec0440000003004000000c23c60c92 0000000168ee32c8b0 00000165888080033ffef5f8454f32cb1bb4203f854dd69bc2ca91b2bce1fb3527440000030000030000030000030050999841d1afd324aea000000300000f600011c0001b40004e40011f0003b80010800059000238000be0005e000220001100000300000300000300000300000300000300000300000300000300000300000300000300000300000300000300000300000300000300000300000300000300000300000300000300000300000300000300000300000300000300000300000300000300000300000300000300000300000300004041 00000001"
	b := EncodeToAVCC(decode(s))
	require.True(t, bytes.HasSuffix(b, []byte{0x40, 0x41}))
	n := naluTypes(b)
	require.Equal(t, []byte{0x67, 0x68, 0x65}, n)
}

func TestFFmpegMPEGTSH264(t *testing.T) {
	// ffmpeg -re -i bbb.mp4 -c copy -f mpegts -
	s := "00000001 09f0 000000016764001fac2484014016ec0440000003004000000c23c60c92 0000000168ee32c8b0 00000165888080033ffef5f8454f32cb1bb4203f854dd69bc2ca91b2bce1fb3527440000030000030000030000030050999841d1afd324aea000000300000f600011c0001b40004e40011f0003b80010800059000238000be0005e000220001100000300000300000300000300000300000300000300000300000300000300000300000300000300000300000300000300000300000300000300000300000300000300000300000300000300000300000300000300000300000300000300000300000300000300000300000300000300000300004041"
	b := EncodeToAVCC(decode(s))
	n := naluTypes(b)
	require.Equal(t, []byte{0x67, 0x68, 0x65}, n)
}

func TestFFmpegHEVC(t *testing.T) {
	// ffmpeg -re -i bbb.mp4 -c libx265 -preset superfast -tune zerolatency -f hevc -
	s := "0000000140010c01ffff01600000030090000003000003005dba0240 0000000142010101600000030090000003000003005da00280802d165bab930bc05a7080000003008000000c04 000000014401c1718312 0000014e0105ffffffffffffffffff2b2ca2de09b51747dbbb55a4fe7fc2fc4e7832363520286275696c642032303829202d20332e352b3131312d6330616631386464353a5b57696e646f77735d5b4743432031332e322e305d5b3634206269745d20386269742b31306269742b3132626974202d20482e3236352f4845564320636f646563202d20436f7079726967687420323031332d3230313820286329204d756c7469636f7265776172"
	b := EncodeToAVCC(decode(s))
	n := naluTypes(b)
	require.Equal(t, []byte{0x40, 0x42, 0x44, 0x4e}, n)
}

func TestFFmpegHEVC2(t *testing.T) {
	// ffmpeg -re -i bbb.mp4 -c libx265 -preset superfast -tune zerolatency -f hevc -
	s := "0000000140010c01ffff01600000030090000003000003005dba0240 0000000142010101600000030090000003000003005da00280802d165bab930bc05a7080000003008000000c04 000000014401c1718312 0000014e0105ffffffffffffffffff2b2ca2de09b51747dbbb55a4fe7fc2fc4e7832363520286275696c642032303829202d20332e352b3131312d6330616631386464353a5b57696e646f77735d5b4743432031332e322e305d5b3634206269745d20386269742b31306269742b3132626974202d20482e3236352f4845564320636f646563202d20436f7079726967687420323031332d3230313820286329204d756c7469636f7265776172652c20496e63202d20687474703a2f2f783236352e6f7267202d206f7074696f6e733a2063707569643d31303439303731206672616d652d746872656164733d32206e756d612d706f6f6c733d3420777070206e6f2d706d6f6465206e6f2d706d65206e6f2d70736e72206e6f2d7373696d206c6f672d6c6576656c3d322062697464657074683d3820696e7075742d6373703d31206670733d32342f3120696e7075742d7265733d313238307837323020696e7465726c6163653d3020746f74616c2d6672616d65733d30206c6576656c2d6964633d3020686967682d746965723d31207568642d62643d30207265663d31206e6f2d616c6c6f772d6e6f6e2d636f6e666f726d616e6365207265706561742d6865616465727320616e6e657862206e6f2d617564206e6f2d656f62206e6f2d656f73206e6f2d68726420696e666f20686173683d302074656d706f72616c2d6c61796572733d30206f70656e2d676f70206d696e2d6b6579696e743d3234206b6579696e743d32353020676f702d6c6f6f6b61686561643d3020626672616d65733d3020622d61646170743d30206e6f2d622d707972616d696420626672616d652d626961733d302072632d6c6f6f6b61686561643d30206c6f6f6b61686561642d736c696365733d34207363656e656375743d30206e6f2d686973742d7363656e65637574207261646c3d30206e6f2d73706c696365206e6f2d696e7472612d72656672657368206374753d3332206d696e2d63752d73697a653d38206e6f2d72656374206e6f2d616d70206d61782d74752d73697a653d33322074752d696e7465722d64657074683d312074752d696e7472612d64657074683d31206c696d69742d74753d302072646f712d6c6576656c3d302064796e616d69632d72643d302e3030206e6f2d7373696d2d7264207369676e68696465206e6f2d74736b6970206e722d696e7472613d30206e722d696e7465723d30206e6f2d636f6e73747261696e65642d696e747261207374726f6e672d696e7472612d736d6f6f7468696e67206d61782d6d657267653d32206c696d69742d726566733d30206e6f2d6c696d69742d6d6f646573206d653d31207375626d653d31206d6572616e67653d35372074656d706f72616c2d6d7670206e6f2d6672616d652d647570206e6f2d686d65206e6f2d77656967687470206e6f2d77656967687462206e6f2d616e616c797a652d7372632d70696373206465626c6f636b3d303a30206e6f2d73616f206e6f2d73616f2d6e6f6e2d6465626c6f636b2072643d322073656c6563746976652d73616f3d30206561726c792d736b69702072736b697020666173742d696e747261206e6f2d74736b69702d66617374206e6f2d63752d6c6f73736c657373206e6f2d622d696e747261206e6f2d73706c697472642d736b697020726470656e616c74793d30207073792d72643d322e3030207073792d72646f713d302e3030206e6f2d72642d726566696e65206e6f2d6c6f73736c65737320636271706f6666733d3020637271706f6666733d302072633d637266206372663d32382e302071636f6d703d302e3630207170737465703d342073746174732d77726974653d302073746174732d726561643d30206970726174696f3d312e34302061712d6d6f64653d302061712d737472656e6774683d302e3030206e6f2d637574726565207a6f6e652d636f756e743d30206e6f2d7374726963742d6362722071672d73697a653d3332206e6f2d72632d677261696e2071706d61783d36392071706d696e3d30206e6f2d636f6e73742d766276207361723d31206f7665727363616e3d3020766964656f666f726d61743d352072616e67653d3020636f6c6f727072696d3d32207472616e736665723d3220636f6c6f726d61747269783d32206368726f6d616c6f633d31206368726f6d616c6f632d746f703d30206368726f6d616c6f632d626f74746f6d3d3020646973706c61792d77696e646f773d3020636c6c3d302c30206d696e2d6c756d613d30206d61782d6c756d613d323535206c6f67322d6d61782d706f632d6c73623d38207675692d74696d696e672d696e666f207675692d6872642d696e666f20736c696365733d31206e6f2d6f70742d71702d707073206e6f2d6f70742d7265662d6c6973742d6c656e6774682d707073206e6f2d6d756c74692d706173732d6f70742d727073207363656e656375742d626961733d302e3035206e6f2d6f70742d63752d64656c74612d7170206e6f2d61712d6d6f74696f6e206e6f2d6864723130206e6f2d68647231302d6f7074206e6f2d6468647231302d6f7074206e6f2d6964722d7265636f766572792d73656920616e616c797369732d72657573652d6c6576656c3d3020616e616c797369732d736176652d72657573652d6c6576656c3d3020616e616c797369732d6c6f61642d72657573652d6c6576656c3d30207363616c652d666163746f723d3020726566696e652d696e7472613d3020726566696e652d696e7465723d3020726566696e652d6d763d3120726566696e652d6374752d646973746f7274696f6e3d30206e6f2d6c696d69742d73616f206374752d696e666f3d30206e6f2d6c6f77706173732d64637420726566696e652d616e616c797369732d747970653d3020636f70792d7069633d31206d61782d617573697a652d666163746f723d312e30206e6f2d64796e616d69632d726566696e65206e6f2d73696e676c652d736569206e6f2d686576632d6171206e6f2d737674206e6f2d6669656c642071702d61646170746174696f6e2d72616e67653d312e3030207363656e656375742d61776172652d71703d30636f6e666f726d616e63652d77696e646f772d6f6666736574732072696768743d3020626f74746f6d3d30206465636f6465722d6d61782d726174653d30206e6f2d7662762d6c6976652d6d756c74692d70617373206e6f2d6d63737466206e6f2d7362726380 0000012801adc2e5bca307b9ce6b18b5ad6a525294a6d117ffd3917322eebaeda718a0000003000003000003021207706824da718a00000300000300000300044408d5db4e31400000030000030000030012500c2725a000000300000300000300002a600e4880000003000003000003000019301180000003000003000003000007d400000300000300000300000300010b000003000003000003000003001810000003000003000003000003019100000300000300000300000d38000003000003000003000067c000000300000300000300025e000003000003000003000c58000003000003000003002b60000003000003000003007f80000003000003000003016300000300000300000303b2000003000003000006e400000300000300000e18000003000003000018d00000030000030000292000000300000300003ce00000030000030000030000030000030000bb80"
	b := EncodeToAVCC(decode(s))
	n := naluTypes(b)
	require.Equal(t, []byte{0x40, 0x42, 0x44, 0x4e, 0x28}, n)
}

func TestFFmpegMPEGTSHEVC(t *testing.T) {
	// ffmpeg -re -i bbb.mp4 -c libx265 -preset superfast -tune zerolatency -an -f mpegts -
	s := "00000001460150 0000000140010c01ffff01600000030090000003000003005dba0240 0000000142010101600000030090000003000003005da00280802d165bab930bc05a7080000003008000000c04 000000014401c1718312 0000014e0105ffffffffffffffffff2b2ca2de09b51747dbbb55a4fe7fc2fc4e7832363520286275696c642032303829202d20332e352b3131312d6330616631386464353a5b57696e646f77735d5b4743432031332e322e305d5b3634206269745d20386269742b31306269742b3132626974202d20482e3236352f4845564320636f646563202d20436f7079726967687420323031332d3230313820286329204d756c7469636f7265776172652c20496e63202d20687474703a2f2f783236352e6f7267202d206f7074696f6e733a2063707569643d31303439303731206672616d652d746872656164733d32206e756d612d706f6f6c733d3420777070206e6f2d706d6f6465206e6f2d706d65206e6f2d70736e72206e6f2d7373696d206c6f672d6c6576656c3d322062697464657074683d3820696e7075742d6373703d31206670733d32342f3120696e7075742d7265733d313238307837323020696e7465726c6163653d3020746f74616c2d6672616d65733d30206c6576656c2d6964633d3020686967682d746965723d31207568642d62643d30207265663d31206e6f2d616c6c6f772d6e6f6e2d636f6e666f726d616e6365207265706561742d6865616465727320616e6e657862206e6f2d617564206e6f2d656f62206e6f2d656f73206e6f2d68726420696e666f20686173683d302074656d706f72616c2d6c61796572733d30206f70656e2d676f70206d696e2d6b6579696e743d3234206b6579696e743d32353020676f702d6c6f6f6b61686561643d3020626672616d65733d3020622d61646170743d30206e6f2d622d707972616d696420626672616d652d626961733d302072632d6c6f6f6b61686561643d30206c6f6f6b61686561642d736c696365733d34207363656e656375743d30206e6f2d686973742d7363656e65637574207261646c3d30206e6f2d73706c696365206e6f2d696e7472612d72656672657368206374753d3332206d696e2d63752d73697a653d38206e6f2d72656374206e6f2d616d70206d61782d74752d73697a653d33322074752d696e7465722d64657074683d312074752d696e7472612d64657074683d31206c696d69742d74753d302072646f712d6c6576656c3d302064796e616d69632d72643d302e3030206e6f2d7373696d2d7264207369676e68696465206e6f2d74736b6970206e722d696e7472613d30206e722d696e7465723d30206e6f2d636f6e73747261696e65642d696e747261207374726f6e672d696e7472612d736d6f6f7468696e67206d61782d6d657267653d32206c696d69742d726566733d30206e6f2d6c696d69742d6d6f646573206d653d31207375626d653d31206d6572616e67653d35372074656d706f72616c2d6d7670206e6f2d6672616d652d647570206e6f2d686d65206e6f2d77656967687470206e6f2d77656967687462206e6f2d616e616c797a652d7372632d70696373206465626c6f636b3d303a30206e6f2d73616f206e6f2d73616f2d6e6f6e2d6465626c6f636b2072643d322073656c6563746976652d73616f3d30206561726c792d736b69702072736b697020666173742d696e747261206e6f2d74736b69702d66617374206e6f2d63752d6c6f73736c657373206e6f2d622d696e747261206e6f2d73706c697472642d736b697020726470656e616c74793d30207073792d72643d322e3030207073792d72646f713d302e3030206e6f2d72642d726566696e65206e6f2d6c6f73736c65737320636271706f6666733d3020637271706f6666733d302072633d637266206372663d32382e302071636f6d703d302e3630207170737465703d342073746174732d77726974653d302073746174732d726561643d30206970726174696f3d312e34302061712d6d6f64653d302061712d737472656e6774683d302e3030206e6f2d637574726565207a6f6e652d636f756e743d30206e6f2d7374726963742d6362722071672d73697a653d3332206e6f2d72632d677261696e2071706d61783d36392071706d696e3d30206e6f2d636f6e73742d766276207361723d31206f7665727363616e3d3020766964656f666f726d61743d352072616e67653d3020636f6c6f727072696d3d32207472616e736665723d3220636f6c6f726d61747269783d32206368726f6d616c6f633d31206368726f6d616c6f632d746f703d30206368726f6d616c6f632d626f74746f6d3d3020646973706c61792d77696e646f773d3020636c6c3d302c30206d696e2d6c756d613d30206d61782d6c756d613d323535206c6f67322d6d61782d706f632d6c73623d38207675692d74696d696e672d696e666f207675692d6872642d696e666f20736c696365733d31206e6f2d6f70742d71702d707073206e6f2d6f70742d7265662d6c6973742d6c656e6774682d707073206e6f2d6d756c74692d706173732d6f70742d727073207363656e656375742d626961733d302e3035206e6f2d6f70742d63752d64656c74612d7170206e6f2d61712d6d6f74696f6e206e6f2d6864723130206e6f2d68647231302d6f7074206e6f2d6468647231302d6f7074206e6f2d6964722d7265636f766572792d73656920616e616c797369732d72657573652d6c6576656c3d3020616e616c797369732d736176652d72657573652d6c6576656c3d3020616e616c797369732d6c6f61642d72657573652d6c6576656c3d30207363616c652d666163746f723d3020726566696e652d696e7472613d3020726566696e652d696e7465723d3020726566696e652d6d763d3120726566696e652d6374752d646973746f7274696f6e3d30206e6f2d6c696d69742d73616f206374752d696e666f3d30206e6f2d6c6f77706173732d64637420726566696e652d616e616c797369732d747970653d3020636f70792d7069633d31206d61782d617573697a652d666163746f723d312e30206e6f2d64796e616d69632d726566696e65206e6f2d73696e676c652d736569206e6f2d686576632d6171206e6f2d737674206e6f2d6669656c642071702d61646170746174696f6e2d72616e67653d312e3030207363656e656375742d61776172652d71703d30636f6e666f726d616e63652d77696e646f772d6f6666736574732072696768743d3020626f74746f6d3d30206465636f6465722d6d61782d726174653d30206e6f2d7662762d6c6976652d6d756c74692d70617373206e6f2d6d63737466206e6f2d7362726380 0000012801adc2e5bca307b9ce6b18b5ad6a525294a6d117ffd3917322eebaeda718a0000003000003000003021207706824da718a00000300000300000300044408d5db4e31400000030000030000030012500c2725a000000300000300000300002a600e4880000003000003000003000019301180000003000003000003000007d400000300000300000300000300010b000003000003000003000003001810000003000003000003000003019100000300000300000300000d38000003000003000003000067c000000300000300000300025e000003000003000003000c58000003000003000003002b60000003000003000003007f80000003000003000003016300000300000300000303b2000003000003000006e400000300000300000e18000003000003000018d00000030000030000292000000300000300003ce00000030000030000030000030000030000bb80"
	b := EncodeToAVCC(decode(s))
	n := naluTypes(b)
	require.Equal(t, []byte{0x40, 0x42, 0x44, 0x4e, 0x28}, n)
}

func TestReolink(t *testing.T) {
	s := "000001460150 00000140010C01FFFF01600000030000030000030000030096AC09 0000000142010101600000030000030000030000030096A001E020021C7F8AAD3BA24BB804000013D800018CE008 000000014401C072F0941E3648 000000012601"
	b := EncodeToAVCC(decode(s))
	n := naluTypes(b)
	require.Equal(t, []byte{0x40, 0x42, 0x44, 0x26}, n)
}

func TestDahua(t *testing.T) {
	s := "00000001460150 00000140010c01ffff01400000030000030000030000030099ac0900 0000000142010101400000030000030000030000030099a001402005a1fe5aee46c1ae550400 000000014401c073c04c9000 000000012601"
	b := EncodeToAVCC(decode(s))
	n := naluTypes(b)
	require.Equal(t, []byte{0x40, 0x42, 0x44, 0x26}, n)
}

func TestUSB(t *testing.T) {
	s := "00 00 00 01 67 4D 00 1F 8D 8D 40 28 02 DD 37 01 01 01 40 00 01 C2 00 00 57 E4 01 00 00 00 01 68 EE 3C 80 00 00 00 00 00 00 00 00 00 00 00 00 00 00 00 00 00 00 00 00 00 00 00 00 00 00 00 01 65 88 80 00"
	b := EncodeToAVCC(decode(s))
	n := naluTypes(b)
	require.Equal(t, []byte{0x67, 0x68, 0x65}, n)

	s = "00 00 00 00 00 00 00 00 00 00 00 00 00 00 01 41 9A 00 4C"
	b = EncodeToAVCC(decode(s))
	n = naluTypes(b)
	require.Equal(t, []byte{0x41}, n)
}
